from torch.utils import data
import numpy as np
import torchvision
import torch

from Transforms import build_transform

__all__ = [
    'build_dataset_CIFAR10',
    'build_dataset_STL10',
    'build_dataset_ImageNet10',
    'build_dataset_ImageNetDogs',
    'build_dataset_CIFAR100',
]

def get_dataloader(train_dataset, test_dataset, args):
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        drop_last=True,
        num_workers=args.workers,
    )

    test_dataloader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=args.test_batch_size,
        shuffle=False,
        drop_last=False,
        num_workers=args.workers,
    )
    args.test_size = len(test_dataset)
    return train_dataloader, test_dataloader


def build_dataset_CIFAR10(dataset_path, args):
    transform_train = build_transform(True, args)
    transform_test = build_transform(False, args)

    train_dataset = torchvision.datasets.CIFAR10(
        download=True,
        root=dataset_path,
        train=True,
        transform=transform_train
    )

    test_dataset = torchvision.datasets.CIFAR10(
        download=True,
        root=dataset_path,
        train=False,
        transform=transform_test
    )
    
    train_dataloader, test_dataloader = get_dataloader(train_dataset, test_dataset, args)
    return train_dataloader, test_dataloader


def build_dataset_STL10(dataset_path, args):
    transform_train = build_transform(True, args)
    transform_test = build_transform(False, args)

    train_dataset = torchvision.datasets.STL10(
        root=dataset_path,
        split='train',
        transform=transform_train
    )

    test_dataset = torchvision.datasets.STL10(
        root=dataset_path,
        split='test',
        transform=transform_test
    )
    
    train_dataloader, test_dataloader = get_dataloader(train_dataset, test_dataset, args)
    return train_dataloader, test_dataloader


def build_dataset_ImageNet10(dataset_path, args):
    transform_train = build_transform(True, args)
    transform_test = build_transform(False, args)

    train_dataset = torchvision.datasets.ImageFolder(
        root=dataset_path,
        transform=transform_train,
    )

    test_dataset = torchvision.datasets.ImageFolder(
        root=dataset_path,
        transform=transform_test,
    )

    train_dataloader, test_dataloader = get_dataloader(train_dataset, test_dataset, args)
    return train_dataloader, test_dataloader


def build_dataset_ImageNetDogs(dataset_path, args):
    transform_train = build_transform(True, args)
    transform_test = build_transform(False, args)

    train_dataset = torchvision.datasets.ImageFolder(
        root=dataset_path,
        transform=transform_train,
    )

    test_dataset = torchvision.datasets.ImageFolder(
        root=dataset_path,
        transform=transform_test,
    )

    train_dataloader, test_dataloader = get_dataloader(train_dataset, test_dataset, args)
    return train_dataloader, test_dataloader

def build_dataset_CIFAR100(dataset_path, args):
    transform_train = build_transform(True, args)
    transform_test = build_transform(False, args)

    if args.concat == False:
        train_dataset = torchvision.datasets.CIFAR100(
            download=True,
            root=dataset_path,
            train=True,
            transform=transform_train
        )

        test_dataset = torchvision.datasets.CIFAR100(
            download=True,
            root=dataset_path,
            train=False,
            transform=transform_test
        )
    else:
        train_dataset_1 = torchvision.datasets.CIFAR100(
            download=True,
            root=dataset_path,
            train=True,
            transform=transform_train
        )
        train_dataset_2 = torchvision.datasets.CIFAR100(
            download=True,
            root=dataset_path,
            train=False,
            transform=transform_train
        )

        test_dataset_1 = torchvision.datasets.CIFAR100(
            download=True,
            root=dataset_path,
            train=True,
            transform=transform_test
        )
        test_dataset_2 = torchvision.datasets.CIFAR100(
            download=True,
            root=dataset_path,
            train=False,
            transform=transform_test
        )
        train_dataset = data.ConcatDataset([train_dataset_1, train_dataset_2])
        test_dataset = data.ConcatDataset([test_dataset_1, test_dataset_2])  

    train_dataset_target = np.array(train_dataset.targets)
    for idx, target in enumerate(train_dataset_target):
        train_dataset_target[idx] = _cifar100_to_cifar20(target)
    train_dataset.targets = train_dataset_target

    test_dataset_target = np.array(test_dataset.targets)
    for idx, target in enumerate(test_dataset_target):
        test_dataset_target[idx] = _cifar100_to_cifar20(target)
    test_dataset.targets = test_dataset_target

    train_dataloader, test_dataloader = get_dataloader(train_dataset, test_dataset, args)
    return train_dataloader, test_dataloader


def _cifar100_to_cifar20(target):
  _dict = \
    {0: 4,
     1: 1,
     2: 14,
     3: 8,
     4: 0,
     5: 6,
     6: 7,
     7: 7,
     8: 18,
     9: 3,
     10: 3,
     11: 14,
     12: 9,
     13: 18,
     14: 7,
     15: 11,
     16: 3,
     17: 9,
     18: 7,
     19: 11,
     20: 6,
     21: 11,
     22: 5,
     23: 10,
     24: 7,
     25: 6,
     26: 13,
     27: 15,
     28: 3,
     29: 15,
     30: 0,
     31: 11,
     32: 1,
     33: 10,
     34: 12,
     35: 14,
     36: 16,
     37: 9,
     38: 11,
     39: 5,
     40: 5,
     41: 19,
     42: 8,
     43: 8,
     44: 15,
     45: 13,
     46: 14,
     47: 17,
     48: 18,
     49: 10,
     50: 16,
     51: 4,
     52: 17,
     53: 4,
     54: 2,
     55: 0,
     56: 17,
     57: 4,
     58: 18,
     59: 17,
     60: 10,
     61: 3,
     62: 2,
     63: 12,
     64: 12,
     65: 16,
     66: 12,
     67: 1,
     68: 9,
     69: 19,
     70: 2,
     71: 10,
     72: 0,
     73: 1,
     74: 16,
     75: 12,
     76: 9,
     77: 13,
     78: 15,
     79: 13,
     80: 16,
     81: 19,
     82: 2,
     83: 4,
     84: 6,
     85: 19,
     86: 5,
     87: 5,
     88: 8,
     89: 19,
     90: 18,
     91: 1,
     92: 2,
     93: 15,
     94: 6,
     95: 0,
     96: 17,
     97: 8,
     98: 14,
     99: 13}

  return _dict[target]